﻿using MathNet.Numerics.LinearAlgebra;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace CoordinateFitterDemo
{
    public class TpsTransform
    {
        private List<PointD> _sourcePoints;
        private double[] _weightsX;
        private double[] _weightsY;
        private double[] _affineCoefficientsX;
        private double[] _affineCoefficientsY;

        /// <summary>
        /// 拟合转换参数
        /// </summary>
        public TpsTransform Fit(IList<PointD> sourcePoints, IList<PointD> targetPoints)
        {
            if (sourcePoints.Count != targetPoints.Count)
                throw new ArgumentException("Source and target points count mismatch");
            if (sourcePoints.Count < 3)
                throw new ArgumentException("At least 3 control points required");

            _sourcePoints = sourcePoints.ToList();
            int n = sourcePoints.Count;

            // 构建TPS矩阵
            Matrix<double> L = Matrix<double>.Build.Dense(n + 3, n + 3);

            // 填充K矩阵（径向基函数部分）
            for (int i = 0; i < n; i++)
            {
                for (int j = 0; j < n; j++)
                {
                    double r2 = SquaredDistance(sourcePoints[i], sourcePoints[j]);
                    L[i, j] = RadialBasisFunction(r2);
                }
            }

            // 填充P矩阵（仿射变换部分）
            for (int i = 0; i < n; i++)
            {
                L[i, n] = 1;
                L[i, n + 1] = sourcePoints[i].X;
                L[i, n + 2] = sourcePoints[i].Y;

                L[n, i] = 1;
                L[n + 1, i] = sourcePoints[i].X;
                L[n + 2, i] = sourcePoints[i].Y;
            }

            // 解线性方程组
            var Vx = Vector<double>.Build.Dense(n + 3);
            var Vy = Vector<double>.Build.Dense(n + 3);

            for (int i = 0; i < n; i++)
            {
                Vx[i] = targetPoints[i].X;
                Vy[i] = targetPoints[i].Y;
            }

            // 使用QR分解求解
            var solver = L.QR().Solve(Vx);
            _weightsX = solver.Take(n).ToArray();
            _affineCoefficientsX = solver.Skip(n).Take(3).ToArray();

            solver = L.QR().Solve(Vy);
            _weightsY = solver.Take(n).ToArray();
            _affineCoefficientsY = solver.Skip(n).Take(3).ToArray();


            return this;
        }

        /// <summary>
        /// 坐标转换
        /// </summary>
        public PointD Transform(PointD point)
        {
            double x = _affineCoefficientsX[0] +
                      _affineCoefficientsX[1] * point.X +
                      _affineCoefficientsX[2] * point.Y;

            double y = _affineCoefficientsY[0] +
                      _affineCoefficientsY[1] * point.X +
                      _affineCoefficientsY[2] * point.Y;

            for (int i = 0; i < _sourcePoints.Count; i++)
            {
                double r2 = SquaredDistance(point, _sourcePoints[i]);
                x += _weightsX[i] * RadialBasisFunction(r2);
                y += _weightsY[i] * RadialBasisFunction(r2);
            }

            return new PointD(x, y);
        }

        private static double RadialBasisFunction(double rSquared)
        {
            return rSquared > 0 ? rSquared * Math.Log(rSquared) : 0;
        }

        private static double SquaredDistance(PointD a, PointD b)
        {
            double dx = a.X - b.X;
            double dy = a.Y - b.Y;
            return dx * dx + dy * dy;
        }
    }

    /// <summary>
    /// 二维坐标点结构
    /// </summary>
    public struct PointD
    {
        public double X { get; }
        public double Y { get; }

        public PointD(double x, double y)
        {
            X = x;
            Y = y;
        }

        public override string ToString() => $"({X:F2}, {Y:F2})";
    }
}
